#include <png.h>
#include <stdio.h>
#include <malloc.h>
#include <unistd.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <libgen.h>
#include <sys/stat.h>

/***
    to png
    指定输入多个文件名，指定输出目录

    restore
    指定输入多个png文件，指定输出目录
**/
/***
    png头：   trix_header 
    - uuid     u8[16]  事务标识
    - number   u32     当前png序号
    - size     u32     当前png有效字节数(不包括头的24字节)

    序号0的png包含总表
    总表：    trix_table
    - name     u32     trix
    - count    u32     文件总数
    - length   u32     后序文件信息表的字节数

    文件信息： trix_file
    - file_size      u32    文件长度
    - file_name_len  u32    文件名长度
    - file_name      string 文件名
**/
static
struct
{
  char   **files;
  char     outdir[4096];
  char     png_prx[64];
  size_t   file_count;
  size_t   max_size;
  int8_t   verbose;
  int8_t   restore;
  int8_t   list;
} options;


#define log_print(fmt, ...) (fprintf(stderr, fmt "\n", ##__VA_ARGS__))
#define log_error(fmt, ...) log_print("ERROR[%d] " fmt, __LINE__, ##__VA_ARGS__)
#define log_info(fmt, ...)                              \
  do {                                                  \
    if (options.verbose) log_print(fmt, ##__VA_ARGS__); \
  } while (0)

#define log_verbose(fmt, ...)                                   \
  do {                                                          \
    if (options.verbose > 1) log_print(fmt, ##__VA_ARGS__);     \
  } while (0)


static
void
arg_error()
{
  log_error("arguments, option -h for help");
  exit(1);
}

static
void
usage()
{
  log_print("usage: [-h][-v][-l][-k MaxSize][-o OutDir][-p OutPrefix] <File>...");
  exit(0);
}

static
char *
uuid2str(const uint8_t *uuid, char *buf)
{
  static const int8_t idx[] =
    {
      3, 2, 1, 0, -1,
      5, 4, -1,
      7, 6, -1,
      8, 9, -1,
      10, 11, 12, 13, 14, 15
    };

  char *p = buf;
  for (size_t i = 0; i != sizeof(idx)/sizeof(idx[0]); ++i)
    {
      if (idx[i] < 0)
        *p++ = '-';
      else
        {
          sprintf(p, "%02x", uuid[idx[i]]);
          p += 2;
        }
    }
  *p = 0;
  return buf;
}

static
int
str2uuid(const char *str, uint8_t *uuid)
{
  unsigned u1, u2, u3, u4;
  long unsigned u5;
  if (sscanf(str, "%x-%x-%x-%x-%lx",
             &u1, &u2, &u3, &u4, &u5) != 5)
    return -1;

  memcpy(uuid + 0, &u1, 4);
  memcpy(uuid + 4, &u2, 2);
  memcpy(uuid + 6, &u3, 2);
  uuid[8] = ((uint8_t *)&u4)[1];
  uuid[9] = ((uint8_t *)&u4)[0];

  uuid[10] = ((uint8_t *)&u5)[5];
  uuid[11] = ((uint8_t *)&u5)[4];
  uuid[12] = ((uint8_t *)&u5)[3];
  uuid[13] = ((uint8_t *)&u5)[2];
  uuid[14] = ((uint8_t *)&u5)[1];
  uuid[15] = ((uint8_t *)&u5)[0];
  return 0;
}

static
int
gen_uuid(uint8_t *uuid)
{
  FILE *fp = fopen("/proc/sys/kernel/random/uuid", "r");
  if (!fp)
    return -1;
  
  char buf[128];
  if (fgets(buf, sizeof(buf), fp) == NULL)
    {
      fclose(fp);
      return -1;
    }
  fclose(fp);
  return str2uuid(buf, uuid);
}

static
int
is_filename_png(const char *fn)
{
  const char *ext = strrchr(fn, '.');
  return (ext && strcasecmp(ext, ".png") == 0);
}

static
void
parse_args(int argc, char **argv)
{
  options.max_size = (512 << 10);
  while (1)
    {
      int opt = getopt(argc, argv, "hvlk:o:p:");
      if (opt == -1)
        break;

      switch (opt)
        {
        case 'h':
          usage();
          break;

        case 'k':
          if (!optarg)
            arg_error();
          
          {
            long unsigned k; 
            if (sscanf(optarg, "%lu", &k) != 1)
              arg_error();
            options.max_size = ((size_t)k << 10);            
          }
          break;

        case 'p':
          if (!optarg)
            arg_error();

          strncpy(options.png_prx, optarg, sizeof(options.png_prx) - 1);
          options.png_prx[sizeof(options.png_prx) - 1] = 0;
          break;

        case 'v':
          options.verbose++ ;
          break;

        case 'l':
          options.list = 1;
          break;

        case 'o':
          if (!realpath(optarg, options.outdir))
            {
              log_error("%s not exits", optarg);
              exit(1);
            }
          break;
          
        case '?':
          break;

        default:
          arg_error();
        }
    }
  options.files = argv + optind;
  if (options.files[0] == NULL) 
    {
      log_error("no files");
      exit(1);
    }
  if (options.outdir[0] == 0 && !realpath(".", options.outdir))
    {
      log_error("no outdir");
      exit(1);
    }
  int is_png = -1;
  for (char **p = options.files; *p; ++p)
    {
      int v = !!is_filename_png(*p);
      if (is_png < 0)
        is_png = v;
      else if (is_png != v)
        {
          log_error("file not match: %s", *p);
          exit(1);
        }
      options.file_count++ ;
    }
  options.restore = is_png;
  if (options.list && !options.restore)
    {
      log_error("option '-l' not right");
    }
}

static int restore();
static int transform();

int
main(int argc, char **argv)
{
  parse_args(argc, argv);

  if (options.restore)
    return !!restore();

  return !!transform();
}

static
int
read_png(FILE *fp, void **pdata, size_t *psize)
{
  png_structp png = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  png_infop info = png_create_info_struct(png);

  if (setjmp(png_jmpbuf(png)))
    { 
      png_destroy_read_struct(&png, &info, NULL);
      return -1;
    }

  png_init_io(png, fp);
  png_read_info(png, info);

  png_byte color_type = png_get_color_type(png, info);
  if (color_type != PNG_COLOR_TYPE_RGBA)
    {
      log_error("color type");
      png_destroy_read_struct(&png, &info, NULL);
      return -1;
    }
  
  int width = png_get_image_width(png, info);
  int height = png_get_image_height(png, info);

  size_t size = (size_t)width * height * 4;
  png_bytep *rows = malloc(height * sizeof(png_bytep));
  if (!rows)
    {
      log_error("malloc");
      png_destroy_read_struct(&png, &info, NULL);
      return -1;
    }

  void *data = malloc(size);
  if (!data)
    {
      free(rows);
      log_error("malloc");
      png_destroy_read_struct(&png, &info, NULL);
      return -1;
    }
  
  for (int y = 0; y < height; y++)
    rows[y] = data + (size_t)y * width * 4;

  png_read_image(png, rows);
  png_destroy_read_struct(&png, &info, NULL);

  free(rows);
  *pdata = data;
  *psize = size;
  return 0;
}

static
long 
int_sqrt(long n)
{
  if (n == 0 || n == 1)
    return n;

  long left = 1, right = n; 
  long result = 0;

  while (left <= right)
    {
      long mid = left + ((right - left) >> 1);
      long square = mid * mid;
      if (square == n)
        return mid;
      else if (square < n)
        {
          result = mid;
          left = mid + 1;
        }
      else
        right = mid - 1;
    }
  return result;
}

static
int
write_png(FILE *fp, void *data, size_t size)
{
  long pixels = ((size + 3) >> 2);
  int height = int_sqrt(pixels);
  int width = pixels / height;

  log_info("size: %zu, pixels: %ld, width: %d, height: %d",
           size, pixels, width, height);
  
  while ((long)width * height < pixels)
    ++ width;

  png_bytep *rows = malloc(sizeof(png_bytep) * height);
  if (!rows)
    {
      log_error("malloc");
      return -1;
    }
    
  for (int n = 0; n < height - 1; ++n)
    rows[n] = data + (size_t)width * 4 * n;

  uint8_t last_line[width * 4];
  memset(last_line, 0, width * 4);
  rows[height - 1] = (void *) last_line;
  size_t last_bytes = size - (size_t)width * (height - 1) * 4;
  memcpy(last_line, data + size - last_bytes, last_bytes);

  png_structp png = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  png_infop info = png_create_info_struct(png);
  if (setjmp(png_jmpbuf(png)))
    {
      free(rows);
      png_destroy_write_struct(&png, &info);
      return -1;
    }
  png_init_io(png, fp);
  png_set_IHDR(png, info, width, height, 8, PNG_COLOR_TYPE_RGBA, 
               PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
  png_write_info(png, info);
  png_write_image(png, rows);
  png_write_end(png, NULL);
  
  png_destroy_write_struct(&png, &info);
  free(rows);
  return 0;
}


struct file_info
{
  size_t  file_size;
  size_t  file_name_len;
  char    file_name[256];
  char    path[4096];

  void   *data;
};

struct trix_header
{
  uint8_t   uuid[16];
  uint32_t  number;
  uint32_t  size;
  uint8_t   data[0];
};

struct trix_table
{
  char      name[4];
  uint32_t  count;
  uint32_t  length;
};

struct trix_file
{
  uint32_t  file_size;
  uint32_t  file_name_len;
  char      file_name[0];
};

struct write_context
{
  uint8_t uuid[16];
  struct trix_header *body;
  size_t    bytes;
  unsigned  number;
};


static
int
trix_write_flush(struct write_context *ctx)
{
  if (ctx->bytes == 0)
    return 0;

  char fname[64 + 32];
  snprintf(fname, sizeof(fname), "%s%04u.png",
           options.png_prx, ctx->number);
  
  char path[4096 + sizeof(fname)];
  snprintf(path, sizeof(path), "%s/%s",
           options.outdir, fname);
  
  FILE *fp = fopen(path, "wb");
  if (!fp)
    {
      log_error("fopen '%s'", path);
      return -1;
    }

  struct trix_header *body = ctx->body;
  body->number = ctx->number;
  body->size = ctx->bytes;
  memcpy(body->uuid, ctx->uuid, 16);

  if (write_png(fp, ctx->body, sizeof(struct trix_header) + ctx->bytes))
    {
      fclose(fp);
      return -1;
    }

  fclose(fp);
  log_print("output %s", fname);
  ctx->bytes = 0;
  ctx->number++ ;
  return 0;
}

static
int
trix_write(struct write_context *ctx, const void *buf, size_t len)
{
  while (len)
    {
      size_t rest = options.max_size - ctx->bytes;
      size_t w = len;
      if (rest < w)
        w = rest;

      memcpy(ctx->body->data + ctx->bytes, buf, w);
      ctx->bytes += w;
      if (ctx->bytes == options.max_size && trix_write_flush(ctx))
        return -1;
      
      buf += w;
      len -= w;
    }
  return 0;
}

static
int
transform()
{
  struct file_info *infos = calloc(sizeof(struct file_info),
                                   options.file_count);
  
  size_t file_bytes = 0; 
  size_t file_entry_len = 0;
  int num = 0;
  for (char **p = options.files; *p; ++p)
    {
      char *fn = *p;
      char  path[4096];
      if (!realpath(fn, path))
        {
          log_error("%s not found", fn);
          return -1;
        }

      struct stat stbuf;
      if (lstat(path, &stbuf) || !S_ISREG(stbuf.st_mode))
        {
          log_error("stat %s", fn);
          return -1;
        }

      char *bn = basename(path);
      size_t bn_len = strlen(bn);
      if (bn_len > 255)
        {
          log_error("basename %s", bn);
          return -1;
        }

      struct file_info *fi = infos + num++;
      fi->file_size = stbuf.st_size;
      fi->file_name_len = bn_len;
      memcpy(fi->file_name, bn, bn_len);
      strncpy(fi->path, path, sizeof(fi->path) - 1);
      fi->path[sizeof(fi->path) - 1] = 0;
      
      file_bytes += fi->file_size;
      file_entry_len += sizeof(struct trix_file) + (((bn_len + 3) >> 2) << 2);
    }

  if (sizeof(struct trix_table) + file_entry_len + 128 > options.max_size)
    options.max_size = sizeof(struct trix_table) + file_entry_len + 128;
  
  struct write_context ctx;
  memset(&ctx, 0, sizeof(ctx));
  gen_uuid(ctx.uuid);
  char uuid_buf[128];
  log_info("uuid: %s", uuid2str(ctx.uuid, uuid_buf));
  
  ctx.body = malloc(options.max_size + sizeof(struct trix_header));
  
  struct trix_table  tab;
  memcpy(tab.name, "trix", 4);
  tab.count = options.file_count;
  tab.length = file_entry_len;
  
  log_info("write trix_table");
  log_info("  file_count: %zu", options.file_count);
  log_info("  length    : %zu", file_entry_len);
  if (trix_write(&ctx, &tab, sizeof(tab)))
    return -1;
  
  for (unsigned n = 0; n != options.file_count; ++n)
    {
      struct file_info *inf = infos + n;
      struct trix_file  entry;
      size_t flen = (((inf->file_name_len + 3) >> 2) << 2);
      entry.file_name_len = flen;
      entry.file_size = inf->file_size;
      
      if (trix_write(&ctx, &entry, sizeof(entry)))
        return -1;
      
      char fn[flen];
      memcpy(fn, inf->file_name, inf->file_name_len);
      memset(fn + inf->file_name_len, 0, flen - inf->file_name_len);
      if (trix_write(&ctx, fn, flen))
        return -1;
    }

  for (unsigned n = 0; n != options.file_count; ++n)
    {
      struct file_info *inf = infos + n;
      
      log_info("write file");
      log_info("  name: %.*s", (int)inf->file_name_len, inf->file_name);
      log_info("  size: %zu", inf->file_size);

      FILE *fp = fopen(inf->path, "rb");
      if (!fp)
        {
          log_error("open file: %s", inf->path);
          return -1;
        }
      for (size_t rest_bytes = inf->file_size; rest_bytes; ) 
        {
          char buf[16*1024];
          size_t w = sizeof(buf);
          if (rest_bytes < w)
            w = rest_bytes;

          if (fread(buf, w, 1, fp) != 1)
            {
              log_error("read file");
              return -1;
            }

          if (trix_write(&ctx, buf, w))
            return -1;
          
          rest_bytes -= w;
        }
      fclose(fp);
    }

  if (trix_write_flush(&ctx))
    return -1;

  return 0;
}

struct entry
{
  char name[256];
  size_t  size;
};

struct restore_context
{
  struct file_info *file, *file_end;
  size_t file_offset;
};

static
int
__make_entry(struct restore_context *ctx,
             const char *fname, size_t fsize)
{
  log_info("make entry %s %zu", fname, fsize);
  
  char path[4096 + 256];
  snprintf(path, sizeof(path), "%s/%s", options.outdir, fname);
  FILE *fp = fopen(path, "wb");
  if (!fp)
    {
      log_error("fopen %s", path);
      return -1;
    }
  while (fsize)
    {
      if (ctx->file == ctx->file_end)
        {
          log_error("data not enough");
          return -1;
        }
      
      size_t rest = ctx->file->file_size - ctx->file_offset;
      size_t bytes = (rest < fsize ? rest : fsize);

      if (fwrite(ctx->file->data + ctx->file_offset, bytes, 1, fp) != 1)
        {
          log_error("fwrite");
          return -1;
        }
      fsize -= bytes;
      ctx->file_offset += bytes;
      if (ctx->file_offset == ctx->file->file_size)
        {
          ++ ctx->file;
          ctx->file_offset = sizeof(struct trix_header);
        }
    }
  fclose(fp);
  return 0;
}

static
int
restore()
{
  struct file_info *infos = calloc(sizeof(struct file_info),
                                   options.file_count);
  size_t sum_size = 0;
  for (char **p = options.files; *p; ++p)
    {
      char *fn = *p;
      char  path[4096];
      if (!realpath(fn, path))
        {
          log_error("%s not found", fn);
          return -1;
        }

      struct stat stbuf;
      if (lstat(path, &stbuf) || !S_ISREG(stbuf.st_mode))
        {
          log_error("stat %s", fn);
          return -1;
        }

      FILE *fp = fopen(path, "rb");
      if (!fp)
        {
          log_error("fopen %s", fn);
          return -1;
        }

      void *data;
      size_t size;
      if (read_png(fp, &data, &size))
        {
          log_error("read png %s", fn);
          return -1;
        }

      if (size < sizeof(struct trix_header))
        {
          log_error("not trix png %s", fn);
          return 1;
        }

      struct trix_header *hdr = data;
      if (hdr->number >= options.file_count)
        {
          log_error("png %s index %u out of range %zu",
                    fn, hdr->number, options.file_count);
          return 1;
        }
      if (hdr->size + sizeof(struct trix_header) > size)
        {
          log_error("png %s size %u large than %zu",
                    fn, hdr->size, size - sizeof(struct trix_header));
          return -1;
        }
            
      struct file_info *inf = infos + hdr->number;
      if (inf->data)
        {
          log_error("png index conflict %s with %s on index %u",
                    fn, inf->file_name, hdr->number);
          return 1;
        }

      inf->data = data;
      /// 仅有效字节，包括文件头
      inf->file_size = hdr->size + sizeof(struct trix_header);
      strncpy(inf->file_name, fn, sizeof(inf->file_name) - 1);
      inf->file_name[sizeof(inf->file_name) - 1] = 0;
      strncpy(inf->path, path, sizeof(inf->path) - 1);
      inf->path[sizeof(inf->path) - 1] = 0;
      fclose(fp);
      sum_size += hdr->size;
    }

  struct trix_header *th = infos->data;
  char uuid_buf[128];
  log_info("uuid: %s", uuid2str(th->uuid, uuid_buf));
  
  for (size_t i = 0; i < options.file_count; ++i)
    {
      struct file_info *inf = infos + i;
      if (!inf->data)
        {
          log_error("png lost #%zu", i);
          return -1;
        }
      struct trix_header *h = inf->data;
      if (i && memcmp(h->uuid, th->uuid, 16))
        {
          log_error("png %s uuid not match", inf->file_name);
          return -1;
        }
      log_info("png #%zu: %s %u", i, inf->file_name, h->size);
    }

  struct trix_table *table = (void *)(th + 1);
  if (memcmp(table->name, "trix", 4))
    {
      log_error("trix table name");
      return -1;
    }

  log_info("table length %u", table->length);
  if (sizeof(struct trix_table) + table->length > th->size)
    {
      log_info("table length");
      return -1;
    }

  log_info("entry count %u", table->count);
  size_t entry_sum_size = 0;
  
  struct entry *objs = calloc(table->count, sizeof(struct entry));

  void *p = table + 1;
  void *p_end = p + table->length;
  for (size_t i = 0; i != table->count; ++i)
    {
      struct trix_file *f = p;
      if ((void *)(f + 1) >= p_end ||
          f->file_name_len == 0 || f->file_name_len > 255)
        {
          log_error("entry exceeds range");
          return -1;
        }

      unsigned fn_len = (((f->file_name_len + 3) >> 2) << 2);
      unsigned entry_len = fn_len + sizeof(struct trix_file);
      if (p + entry_len > p_end)
        {
          log_error("entry exceeds range");
          return -1;
        }

      struct entry *ent = objs + i;
      memcpy(ent->name, f->file_name, fn_len);
      ent->size = f->file_size;

      log_print("entry #%zu: %s %zu", i, ent->name, ent->size);
      
      entry_sum_size += ent->size;
      p += entry_len;
    }

  log_info("all entry size %zu", entry_sum_size);
  if (entry_sum_size + sizeof(struct trix_table) + table->length != sum_size)
    {
      log_error("size mot match");
      return -1;
    }

  if (options.list)
    return 0;

  struct restore_context ctx;
  ctx.file = infos;
  ctx.file_end = infos + options.file_count;
  ctx.file_offset = sizeof(struct trix_header) + sizeof(struct trix_table) + table->length;
  
  for (size_t i = 0; i != table->count; ++i)
    {
      struct entry *ent = objs + i;
      if (__make_entry(&ctx, ent->name, ent->size))
        return -1;
    }

  return 0;
}
